package cas.ihep.util;

/**
 * Created by chocolate on 2017/11/24.
 */

import org.apache.commons.collections.keyvalue.MultiKey;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.*;

import java.io.IOException;
import java.util.*;

public class FileLocation{
    private FileSystem hdfs;
    // mapping from a block to the nodes on which it has replicas
    private static final PathFilter hiddenFileFilter = new PathFilter(){
        public boolean accept(Path p){
            String name = p.getName();
            return !name.startsWith("_") && !name.startsWith(".");
        }
    };
    public FileLocation() throws IOException {
        Configuration config = new Configuration();
        hdfs = FileSystem.get(config);
    }

    public List<FileStatus> listStatus(String dir) throws IOException {
        List<FileStatus> result = new ArrayList<FileStatus>();
        Path path = new Path(dir);
        List<PathFilter> filters = new ArrayList<PathFilter>();
        filters.add(hiddenFileFilter);
        PathFilter inputFilter = new MultiPathFilter(filters);
        FileStatus[] matches = hdfs.globStatus(path, inputFilter);
        for (FileStatus globStat: matches) {
            if (globStat.isDirectory()) {
                for(FileStatus stat: hdfs.listStatus(globStat.getPath(),inputFilter)) {
                    result.add(stat);
                }
            } else {
                result.add(globStat);
            }
        }
        return result;
    }


    public String[] GetFileLocations(FileStatus fileStatus) throws IOException {
        BlockLocation[] blkLocations = hdfs.getFileBlockLocations(fileStatus, 0, fileStatus.getLen());
        return blkLocations[0].getHosts();
    }


    public static void main(String[] args) throws Exception {
        if(args.length <1){
            System.out.println("Usage FileLocation dirname");
            System.exit(1);
        }
        HashMap<MultiKey,List<Path>> nodesToFiles = new HashMap<MultiKey, List<Path>>();
        FileLocation test = new FileLocation();

        List<FileStatus> statuses= test.listStatus(args[0]);
        if (statuses == null)
            System.exit(2);
        System.out.println("get" + statuses.size() + "files!");
        for(FileStatus fileStatus: statuses){
            String[] locations = test.GetFileLocations(fileStatus);
            Arrays.sort(locations);
            MultiKey multiKey = new MultiKey(locations);
            List<Path> paths = nodesToFiles.get(multiKey);
            if( paths == null){
                paths = new ArrayList<Path>();
                nodesToFiles.put(multiKey,paths);
            }
            paths.add(fileStatus.getPath());
        }
        Iterator<Map.Entry<MultiKey,List<Path>>> iter1 = nodesToFiles.entrySet().iterator();
        while (iter1.hasNext()) {
            Map.Entry<MultiKey,List<Path>> entry =  iter1.next();
            MultiKey multiKey = entry.getKey();
            List<Path> paths = entry.getValue();
            System.out.print(StringUtils.join(multiKey.getKeys(),",") + "(" + paths.size() + "): ");
            System.out.print("{");
            int i =0;
            for(Path path:paths){
                if(i>0){
                    System.out.print(",");
                }
                System.out.print("\"" + path.getName()+"\"");
                i++;
            }
            System.out.print("}");
            System.out.println("");
        }
        System.exit(0);
    }
    private static class MultiPathFilter implements PathFilter {
        private List<PathFilter> filters;

        public MultiPathFilter(List<PathFilter> filters) {
            this.filters = filters;
        }

        public boolean accept(Path path) {
            for (PathFilter filter : filters) {
                if (!filter.accept(path)) {
                    return false;
                }
            }
            return true;
        }
    }
}

